import asyncio
from typing import Optional, Dict, Any
from ..backends import BackendManager
from .schemas import (
    ChatCompletionRequest, 
    ChatCompletionResponse,
    CompletionRequest,
    CompletionResponse
)


class ModelManager:
    """Manages model loading, unloading, and inference across different backends."""
    
    def __init__(self):
        self.backend_manager = BackendManager()
        self.current_model = None
        self.current_backend = None
        
    async def load_model(self, model_name: str, backend: str, model_path: Optional[str] = None):
        """Load a model using the specified backend."""
        if self.current_model:
            await self.unload_model()
            
        self.current_backend = self.backend_manager.get_backend(backend)
        self.current_model = await self.current_backend.load_model(model_name, model_path)
        
    async def unload_model(self):
        """Unload the currently loaded model."""
        if self.current_model and self.current_backend:
            await self.current_backend.unload_model(self.current_model)
            self.current_model = None
            self.current_backend = None
            
    async def generate_chat_completion(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
        """Generate a chat completion using the loaded model."""
        if not self.current_model or not self.current_backend:
            raise RuntimeError("No model loaded. Call load_model first.")
            
        return await self.current_backend.generate_chat(self.current_model, request)
        
    async def generate_completion(self, request: CompletionRequest) -> CompletionResponse:
        """Generate a completion using the loaded model."""
        if not self.current_model or not self.current_backend:
            raise RuntimeError("No model loaded. Call load_model first.")
            
        return await self.current_backend.generate_completion(self.current_model, request)